import pandas as pd
import numpy as np
from datetime import datetime, timedelta

def calculate_paying_buyers(
        orders_df: pd.DataFrame,
        pre_sale_orders_df: pd.DataFrame = None,
        start_date: str = None,
        end_date: str = None
) -> int:
    """
    计算指定时间范围内的支付买家数

    参数:
    - orders_df: 订单数据DataFrame
    - pre_sale_orders_df: 预售订单关联数据DataFrame（可选）
    - start_date: 开始日期（格式: 'YYYY-MM-DD'）
    - end_date: 结束日期（格式: 'YYYY-MM-DD'）

    返回:
    - 支付买家数（整数）
    """
    # 过滤时间范围
    if start_date and end_date:
        filtered_orders = orders_df[
            (orders_df['pay_time'] >= start_date) &
            (orders_df['pay_time'] <= end_date)
            ]
    else:
        filtered_orders = orders_df.copy()

    # 处理预售订单逻辑
    if pre_sale_orders_df is not None and not pre_sale_orders_df.empty:
        # 提取预售尾款订单ID
        balance_orders = pre_sale_orders_df['balance_order_id'].dropna()

        # 只保留预售尾款订单和普通订单
        filtered_orders = filtered_orders[
            (filtered_orders['order_id'].isin(balance_orders)) |
            (filtered_orders['order_type'] == 'normal')
            ]

    # 过滤出支付成功的订单
    paid_orders = filtered_orders[
        filtered_orders['pay_status'] == 'paid'
        ]

    # 去重计算买家数
    unique_buyers = paid_orders['buyer_id'].nunique()

    return unique_buyers


# 示例调用
if __name__ == "__main__":
    # 模拟数据（实际使用时替换为真实数据）
    orders_df = pd.DataFrame({
        'order_id': ['ORD001', 'ORD002', 'ORD003', 'ORD004', 'ORD005'],
        'buyer_id': ['U001', 'U001', 'U002', 'U003', 'U003'],
        'order_type': ['normal', 'normal', 'pre_sale', 'normal', 'pre_sale'],
        'pay_status': ['paid', 'paid', 'paid', 'failed', 'paid'],
        'pay_time': [
            '2025-07-01 10:00:00',
            '2025-07-02 14:30:00',
            '2025-07-03 09:15:00',
            '2025-07-04 16:45:00',
            '2025-07-05 11:20:00'
        ],
        'refund_time': [
            '2025-07-03 15:00:00',  # U001的ORD001已退款
            None,
            None,
            None,
            None
        ],
        'payment_method': ['alipay', 'wechat', 'alipay', 'alipay', 'wechat']
    })

    # 计算7月1日至7月5日的支付买家数
    result = calculate_paying_buyers(
        orders_df=orders_df,
        start_date='2025-07-01',
        end_date='2025-07-05'
    )

    print(f"支付买家数 (2025-07-01至2025-07-05): {result}")
    # 输出应为3（U001两次支付但只算1次，U002一次，U003一次）